#!/bin/bash

# Kafka 集群管理脚本
# 功能：启动/停止集群、创建/删除/查看带自定义参数的 Topic，操作后返回菜单

# 配置信息
KAFKA_HOME="/opt/module/kafka_2.13-4.0.0"  # 修改为你的 Kafka 安装路径
SERVER_CONFIG="config/server.properties"
BOOTSTRAP_SERVERS="node1:9092,node2:9092,node3:9092"  # 修改为你的 Kafka 集群地址
LOG_DIR="/opt/logs/脚本"  # 日志目录
LOG_FILE="$LOG_DIR/kafka_manager.log"

# 集群节点列表（用于启停所有节点）
KAFKA_NODES=("node1" "node2" "node3")  # 修改为你的 Kafka 节点列表

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# 创建日志目录
create_log_dir() {
    if [ ! -d "$LOG_DIR" ]; then
        echo -e "${YELLOW}创建日志目录: $LOG_DIR${NC}"
        mkdir -p "$LOG_DIR" || error_exit "无法创建日志目录: $LOG_DIR"
    fi
}

# 日志函数
log() {
    local level=$1
    local message=$2
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    
    # 确保日志目录存在
    create_log_dir
    
    echo "[$timestamp] [$level] $message" | tee -a "$LOG_FILE"
}

# 错误处理函数
error_exit() {
    log "ERROR" "$1"
    echo -e "${RED}错误: $1${NC}" >&2
    exit 1
}

# 检查命令是否存在
check_command() {
    command -v "$1" >/dev/null 2>&1 || error_exit "未找到命令: $1"
}

# 检查 Kafka 安装目录
check_kafka_home() {
    if [ ! -d "$KAFKA_HOME" ]; then
        error_exit "Kafka 安装目录不存在: $KAFKA_HOME"
    fi
}

# 检查远程节点上的 Kafka 服务是否启动
check_remote_kafka() {
    local node=$1
    
    log "INFO" "检查节点 $node 上的 Kafka 进程..."
    
    # 通过 SSH 使用 jps 检查 Kafka 进程
    if ssh "$node" "jps | grep -q 'Kafka$'"; then
        log "INFO" "节点 $node 上的 Kafka 进程正在运行"
        return 0
    else
        log "INFO" "节点 $node 上的 Kafka 进程未运行"
        return 1
    fi
}

# 启动单个 Kafka 节点
start_kafka_node() {
    local node=$1
    
    log "INFO" "在节点 $node 上启动 Kafka..."
    echo -e "${YELLOW}正在节点 $node 上启动 Kafka...${NC}"
    
    # 通过 SSH 启动 Kafka
    ssh "$node" "$KAFKA_HOME/bin/kafka-server-start.sh -daemon $KAFKA_HOME/$SERVER_CONFIG"
    
    # 等待并检查启动状态
    local max_attempts=10
    local attempt=1
    
    while [ $attempt -le $max_attempts ]; do
        if check_remote_kafka "$node"; then
            log "INFO" "节点 $node 上的 Kafka 启动成功"
            echo -e "${GREEN}节点 $node 上的 Kafka 启动成功${NC}"
            return 0
        fi
        
        echo -e "${YELLOW}等待节点 $node 上的 Kafka 进程启动 ($attempt/$max_attempts)...${NC}"
        sleep 3
        attempt=$((attempt + 1))
    done
    
    log "ERROR" "节点 $node 上的 Kafka 启动超时"
    echo -e "${RED}节点 $node 上的 Kafka 启动超时${NC}"
    return 1
}

# 停止单个 Kafka 节点
stop_kafka_node() {
    local node=$1
    
    log "INFO" "在节点 $node 上停止 Kafka..."
    echo -e "${YELLOW}正在节点 $node 上停止 Kafka...${NC}"
    
    # 通过 SSH 停止 Kafka
    ssh "$node" "$KAFKA_HOME/bin/kafka-server-stop.sh"
    
    # 等待并检查停止状态
    local max_attempts=10
    local attempt=1
    
    while [ $attempt -le $max_attempts ]; do
        if ! check_remote_kafka "$node"; then
            log "INFO" "节点 $node 上的 Kafka 已停止"
            echo -e "${GREEN}节点 $node 上的 Kafka 已停止${NC}"
            return 0
        fi
        
        echo -e "${YELLOW}等待节点 $node 上的 Kafka 进程停止 ($attempt/$max_attempts)...${NC}"
        sleep 3
        attempt=$((attempt + 1))
    done
    
    log "WARN" "节点 $node 上的 Kafka 未正常停止，尝试强制终止"
    echo -e "${RED}节点 $node 上的 Kafka 未正常停止，尝试强制终止...${NC}"
    
    # 强制终止
    ssh "$node" "jps | grep 'Kafka$' | awk '{print \$1}' | xargs kill -9"
    sleep 2
    
    if ! check_remote_kafka "$node"; then
        log "INFO" "节点 $node 上的 Kafka 已强制终止"
        echo -e "${GREEN}节点 $node 上的 Kafka 已强制终止${NC}"
        return 0
    else
        log "ERROR" "无法停止节点 $node 上的 Kafka 服务"
        echo -e "${RED}无法停止节点 $node 上的 Kafka 服务${NC}"
        return 1
    fi
}

# 启动 Kafka 集群
start_cluster() {
    log "INFO" "正在启动 Kafka 集群..."
    echo -e "${YELLOW}正在启动 Kafka 集群...${NC}"
    
    local success=true
    
    # 逐个启动节点
    for node in "${KAFKA_NODES[@]}"; do
        if ! start_kafka_node "$node"; then
            success=false
        fi
    done
    
    if [ "$success" = true ]; then
        log "INFO" "Kafka 集群启动成功"
        echo -e "${GREEN}Kafka 集群启动成功${NC}"
    else
        error_exit "Kafka 集群启动失败，部分节点未正常启动"
    fi
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
}

# 停止 Kafka 集群
stop_cluster() {
    log "INFO" "正在停止 Kafka 集群..."
    echo -e "${YELLOW}正在停止 Kafka 集群...${NC}"
    
    local success=true
    
    # 逐个停止节点（逆序）
    for (( idx=${#KAFKA_NODES[@]}-1 ; idx>=0 ; idx-- )); do
        node=${KAFKA_NODES[$idx]}
        if ! stop_kafka_node "$node"; then
            success=false
        fi
    done
    
    if [ "$success" = true ]; then
        log "INFO" "Kafka 集群已停止"
        echo -e "${GREEN}Kafka 集群已停止${NC}"
    else
        error_exit "Kafka 集群停止失败，部分节点未正常停止"
    fi
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
}

# 创建 Topic
create_topic() {
    local topic_name=$1
    local partitions=$2
    local replication_factor=$3
    
    log "INFO" "创建 Topic: $topic_name (分区: $partitions, 副本: $replication_factor)"
    echo -e "${YELLOW}正在创建 Topic: $topic_name${NC}"
    
    # 创建 Topic
    "$KAFKA_HOME/bin/kafka-topics.sh" --create \
        --bootstrap-server "$BOOTSTRAP_SERVERS" \
        --topic "$topic_name" \
        --partitions "$partitions" \
        --replication-factor "$replication_factor"
    
    if [ $? -ne 0 ]; then
        error_exit "创建 Topic 失败: $topic_name"
    fi
    
    # 验证 Topic 创建
    if describe_topic "$topic_name"; then
        log "INFO" "Topic 创建成功: $topic_name"
        echo -e "${GREEN}Topic 创建成功: $topic_name${NC}"
    else
        error_exit "Topic 创建后验证失败: $topic_name"
    fi
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
}

# 查看 Topic 详细信息
describe_topic() {
    local topic_name=$1
    
    log "INFO" "查看 Topic 详细信息: $topic_name"
    echo -e "${YELLOW}Topic 详细信息: $topic_name${NC}"
    
    # 查看 Topic 详细信息
    "$KAFKA_HOME/bin/kafka-topics.sh" --describe \
        --bootstrap-server "$BOOTSTRAP_SERVERS" \
        --topic "$topic_name"
    
    local result=$?
    if [ $result -ne 0 ]; then
        log "ERROR" "查看 Topic 详细信息失败: $topic_name"
        echo -e "${RED}查看 Topic 详细信息失败: $topic_name${NC}"
    fi
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
    return $result
}

# 列出所有 Topics
list_topics() {
    log "INFO" "列出所有 Topics"
    echo -e "${YELLOW}所有 Topics:${NC}"
    
    # 列出所有 Topics
    "$KAFKA_HOME/bin/kafka-topics.sh" --list \
        --bootstrap-server "$BOOTSTRAP_SERVERS"
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
}

# 删除 Topic
delete_topic() {
    local topic_name=$1
    
    # 确认删除操作
    read -p "确定要删除 Topic '$topic_name' 吗？(y/N): " confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        log "INFO" "取消删除 Topic: $topic_name"
        echo -e "${YELLOW}取消删除 Topic: $topic_name${NC}"
        # 操作取消后暂停
        read -p "按 Enter 键返回菜单..."
        return 1
    fi
    
    log "INFO" "删除 Topic: $topic_name"
    echo -e "${YELLOW}正在删除 Topic: $topic_name${NC}"
    
    # 删除 Topic
    "$KAFKA_HOME/bin/kafka-topics.sh" --delete \
        --bootstrap-server "$BOOTSTRAP_SERVERS" \
        --topic "$topic_name"
    
    if [ $? -ne 0 ]; then
        error_exit "删除 Topic 失败: $topic_name"
    fi
    
    log "INFO" "Topic 删除请求已提交: $topic_name"
    echo -e "${GREEN}Topic 删除请求已提交: $topic_name${NC}"
    echo -e "${YELLOW}注意: Kafka 可能需要一些时间来完全删除 Topic${NC}"
    
    # 操作完成后暂停
    read -p "按 Enter 键返回菜单..."
    return 0
}

# 显示帮助信息
show_help() {
    echo "Kafka 集群管理脚本"
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  --start          启动 Kafka 集群"
    echo "  --stop           停止 Kafka 集群"
    echo "  --create-topic   创建新 Topic (需要提供 --topic-name 参数)"
    echo "  --delete-topic   删除 Topic (需要提供 --topic-name 参数)"
    echo "  --describe-topic 查看 Topic 详细信息 (需要提供 --topic-name 参数)"
    echo "  --list-topics    列出所有 Topics"
    echo "  --topic-name     Topic 名称 (与 --create-topic/--delete-topic/--describe-topic 一起使用)"
    echo "  --partitions     分区数 (默认: 3)"
    echo "  --replication    副本因子 (默认: 3)"
    echo "  --help           显示此帮助信息"
    exit 0
}

# 显示主菜单并获取用户选择
show_menu() {
    while true; do
        clear
        echo "=== Kafka 集群管理 ==="
        echo "1. 启动集群"
        echo "2. 停止集群"
        echo "3. 创建 Topic"
        echo "4. 删除 Topic"
        echo "5. 查看 Topic 详细信息"
        echo "6. 列出所有 Topics"
        echo "7. 退出"
        read -p "请选择操作 [1-7]: " choice
        
        case $choice in
            1)
                start_cluster
                ;;
            2)
                stop_cluster
                ;;
            3)
                read -p "输入 Topic 名称: " topic_name
                read -p "输入分区数 (默认 3): " partitions_input
                read -p "输入副本因子 (默认 3): " replication_input
                
                # 设置默认值
                partitions=${partitions_input:-3}
                replication_factor=${replication_input:-3}
                
                create_topic "$topic_name" "$partitions" "$replication_factor"
                ;;
            4)
                read -p "输入要删除的 Topic 名称: " topic_name
                delete_topic "$topic_name"
                ;;
            5)
                read -p "输入要查看的 Topic 名称: " topic_name
                describe_topic "$topic_name"
                ;;
            6)
                list_topics
                ;;
            7)
                echo "退出..."
                exit 0
                ;;
            *)
                echo -e "${RED}无效选择，请重新输入${NC}"
                sleep 1
                ;;
        esac
    done
}

# 主函数
main() {
    check_command "jps"
    check_command "ssh"
    check_kafka_home
    
    # 如果有命令行参数，执行相应操作后退出
    if [[ $# -gt 0 ]]; then
        local action=""
        local topic_name=""
        local partitions=3
        local replication_factor=3
        
        # 解析命令行参数
        while [[ $# -gt 0 ]]; do
            case $1 in
                --start)
                    action="start"
                    ;;
                --stop)
                    action="stop"
                    ;;
                --create-topic)
                    action="create_topic"
                    ;;
                --delete-topic)
                    action="delete_topic"
                    ;;
                --describe-topic)
                    action="describe_topic"
                    ;;
                --list-topics)
                    action="list_topics"
                    ;;
                --topic-name)
                    topic_name=$2
                    shift
                    ;;
                --partitions)
                    partitions=$2
                    shift
                    ;;
                --replication)
                    replication_factor=$2
                    shift
                    ;;
                --help)
                    show_help
                    ;;
                *)
                    echo "未知参数: $1" >&2
                    show_help
                    ;;
            esac
            shift
        done
        
        # 执行命令行指定的操作
        case $action in
            start)
                start_cluster
                ;;
            stop)
                stop_cluster
                ;;
            create_topic)
                if [ -z "$topic_name" ]; then
                    error_exit "创建 Topic 时必须指定 --topic-name 参数"
                fi
                create_topic "$topic_name" "$partitions" "$replication_factor"
                ;;
            delete_topic)
                if [ -z "$topic_name" ]; then
                    error_exit "删除 Topic 时必须指定 --topic-name 参数"
                fi
                delete_topic "$topic_name"
                ;;
            describe_topic)
                if [ -z "$topic_name" ]; then
                    error_exit "查看 Topic 详细信息时必须指定 --topic-name 参数"
                fi
                describe_topic "$topic_name"
                ;;
            list_topics)
                list_topics
                ;;
            *)
                error_exit "未指定有效操作"
                ;;
        esac
        
        exit 0
    fi
    
    # 没有命令行参数时，显示交互式菜单
    show_menu
}

# 执行主函数
main "$@"    